package leecode.topk;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.PriorityQueue;

/**
 * @author wangxi created on 2021/5/16 16:18
 * @version v1.0
 *
 * https://leetcode-cn.com/problems/top-k-frequent-words/
 * 给一非空的单词列表，返回前 k 个出现次数最多的单词。
 *
 * 返回的答案应该按单词出现频率由高到低排序。如果不同的单词有相同出现频率，按字母顺序排序。
 */
public class TopKFrequentWord {

    public static void main(String[] args) {
        TopKFrequentWord obj = new TopKFrequentWord();
        String[] words = {"the", "day", "is", "sunny", "the", "the", "the", "sunny", "is", "is"};
        List<String> list = obj.topKFrequent(words, 4);

        System.out.println(list);
    }

    public List<String> topKFrequent(String[] words, int k) {
        List<String> res = new ArrayList<>();;
        if (words == null || words.length <= 0) {
            return res;
        }
        Map<String, Integer> dict = new HashMap<>();
        for (String s : words) {
            dict.put(s, dict.getOrDefault(s, 0) + 1);
        }
        // 堆的排序是保留哪些元素
        PriorityQueue<String> queue = new PriorityQueue<>((o1, o2) -> {
            if (dict.get(o1).equals(dict.get(o2)) ) {
                // 这里本应该是 o1.compareTo(o2) 小到大排序
                // 但是在后面会把res集合reverse，因此这里先按照字母顺序倒序
                return o2.compareTo(o1);
            } else {
                return Integer.compare(dict.get(o1), dict.get(o2));
            }
        });
        // 注意：使用的是dict.keySet()，使用原来的words数据会有问题
        for (String s : dict.keySet()) {
            queue.offer(s);
            if (queue.size() > k) {
                // 自动弹出最小元素
                queue.poll();
            }
        }
        // 堆可以这样遍历，这种是小到大排序
        while (!queue.isEmpty()) {
            res.add(queue.poll());
        }
        // topk -> 大到小
        Collections.reverse(res);
        return res;
    }
}
